package org.easy.byteaidraw;


import org.easy.byteaidraw.dto.AnimeDrawRequest;
import org.easy.core.dto.ImgSpec;
import org.easy.core.util.Lists;
import org.easy.byteaidraw.dto.RoleFeatureRetentionRequest;
import lombok.extern.slf4j.Slf4j;

import java.util.List;


/**
 * 通用方法封装工具类
 */
@Slf4j
public class ByteAiDrawTemplate {

    private ByteAiDrawClient byteAiDrawClient;


    public ByteAiDrawTemplate(ByteAiDrawClient client) {
        this.byteAiDrawClient = client;
    }

    public String animeDraw(AnimeDrawRequest animeDrawRequest) {
        String string = byteAiDrawClient.animeDrawSubmit(animeDrawRequest);
        log.info("生图提交成功，等待查询taskId{}", string);
        return waitForResult(string);
    }

    public String animeDraw(String prompt, ImgSpec imgSpec) {
        AnimeDrawRequest animeDrawRequest = new AnimeDrawRequest();
        animeDrawRequest.setPrompt(prompt);
        animeDrawRequest.setWidth(imgSpec.getWidth());
        animeDrawRequest.setHeight(imgSpec.getHeight());
        String taskId = byteAiDrawClient.animeDrawSubmit(animeDrawRequest);
        log.info("生图提交成功，等待查询taskId{}", taskId);
        return waitForResult(taskId);
    }

    public String animeDraw(String prompt, ImgSpec imgSpec, String sourceImage) {
        AnimeDrawRequest animeDrawRequest = new AnimeDrawRequest();
        animeDrawRequest.setPrompt(prompt);
        animeDrawRequest.setOriginImg(sourceImage);
        animeDrawRequest.setWidth(imgSpec.getWidth());
        animeDrawRequest.setHeight(imgSpec.getHeight());
        String taskId = byteAiDrawClient.animeDrawSubmit(animeDrawRequest);
        log.info("生图提交成功，等待查询taskId{}", taskId);
        return waitForResult(taskId);
    }

    public String waitForResult(String generateUuid) {
        String string = byteAiDrawClient.animeDrawGetResult(generateUuid);
        while (string == null) {
            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            string = byteAiDrawClient.animeDrawGetResult(generateUuid);
        }
        return string;
    }


    public String roleFeatureRetention(RoleFeatureRetentionRequest request) {
        String taskId = byteAiDrawClient.roleFeatureRetention(request);
        String result = byteAiDrawClient.animeDrawGetResult(taskId);
        while (result == null) {
            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            result = byteAiDrawClient.waitForResultByRoleFeatureRetention(taskId);
        }
        return result;
    }

    public String byteedit_v2_0(String prompt, String originImg, Float scale) {
        String taskId = byteAiDrawClient.byteedit_v2_0(prompt, originImg, scale);
        String result = byteAiDrawClient.animeDrawGetResult(taskId);
        while (result == null) {
            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            result = byteAiDrawClient.waitForResultByRoleFeatureRetention(taskId);
        }
        return result;
    }


    public String img2img_anime_accelerated_maintain_id_for_smart_drawing_anime(String prompt, ImgSpec imgSpec, String originImg) {
        List<String> originImgs = Lists.newArrayList(originImg, originImg);
        String generateUuid = byteAiDrawClient.img2img_anime_accelerated_maintain_id_for_smart_drawing_anime(prompt, originImgs, imgSpec.getWidth(), imgSpec.getHeight());
        String result = null;
        while (result == null) {
            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            result = byteAiDrawClient.waitForResultByimg2img_anime_accelerated_maintain_id_for_smart_drawing_anime(generateUuid);
        }
        return result;
    }


}